Summary

Fix #3503: HunyuanImage3 AR output on NPU is missing the <recaption> opening tag after </think>, which breaks the downstream DiT stage. The same code path (prefer_model_sampler) is also used by CosyVoice3, so this fix benefits any model that opts into custom sampling.

Root Cause

HunyuanImage3 declares prefer_model_sampler = True and implements a custom sample() method that ports the official _StageTransitionLogitsProcessor. After </think>, it overrides logits to force <recaption> (and analogous transitions for </recaption>).

The GPU AR runner honors this contract at gpu_ar_model_runner.py::_sample:

if logits is not None and callable(model_sample) and \
        getattr(self.model, "prefer_model_sampler", False):
    sampler_output = model_sample(
        logits,
        self._sampling_metadata_for_model_sampler(sampling_metadata),
    )
    if sampler_output is not None:
        return sampler_output

NPUARModelRunner had no such override. It inherited vllm-ascend.NPUModelRunner._sample, which is unconditionally:

return self.sampler(logits=logits, sampling_metadata=sampling_metadata)

So model.sample() was never called on NPU — the stage transition forcing logic was completely bypassed. With sampling temperature=0.6 / top_p=0.95 / top_k=1024, the model would freely sample whatever after </think>, which on Chinese prompts almost always landed on actual text (请...) instead of the <recaption> token.

This is purely an integration gap, not a sampler-algorithm difference between CUDA and NPU. The dispatch hook was added in #1703 for CosyVoice3 and gates on the generic prefer_model_sampler attribute, so HunyuanImage3 (#2713) opted into it for free on GPU. NPU never picked up the same generalization.

Fix

1. Move _build_model_sampler_output_token_ids and _sampling_metadata_for_model_sampler from GPUARModelRunner to OmniGPUModelRunner. They are pure logic over self.input_batch with no device-specific code.

   The MRO for OmniNPUModelRunner is OmniNPUModelRunner → OmniGPUModelRunner → NPUModelRunner → ..., so NPU inherits the helpers automatically.

2. Add a thin _sample override in NPUARModelRunner mirroring the GPU one. On the fall-through (no model sampler, or spec-decode), call super()._sample(...) so NPU keeps its lmhead_tp_enable logits slicing and rejection_sampler path unchanged.

Touched files (+83 / -48):

• vllm_omni/worker/gpu_model_runner.py — host the two shared helpers
• vllm_omni/worker/gpu_ar_model_runner.py — drop the duplicates
• vllm_omni/platforms/npu/worker/npu_ar_model_runner.py — drop unused imports, add _sample override

Future Cleanup

This PR is the minimal correctness fix. Two follow-ups worth doing once we have more prefer_model_sampler users or another platform:

1. Push _sample itself down to OmniGPUModelRunner. The dispatch logic is identical between GPU and NPU; only the fall-through target differs, and super()._sample(...) resolves correctly via MRO on both sides (GPU → vllm.GPUModelRunner._sample, NPU → vllm-ascend.NPUModelRunner._sample with lmhead_tp_enable slicing and rejection sampler). Holding off here only because the logit-bias call uses self.sampler.logit_bias_state, and we want one more pair of eyes on whether that shape is identical on both platforms before merging the override.

2. Collapse GPUARModelRunner / NPUARModelRunner duplication. The two files are ~1040 lines each with ~95% structural overlap (_request_final_stage_id, _request_needs_downstream_stage_payload, _resolve_pooler_payload_req_ids, _resolve_req_hidden_states, _maybe_update_prefix_cache, _resolve_global_request_id, the propose_draft_token_ids shape, etc.). Two viable shapes:

   • Mixin (ModelSamplerDispatchMixin, etc.): smallest blast radius, both runners mix in. Each new shared concern becomes its own mixin.
   • Sibling-merge via NPUARModelRunner(GPUARModelRunner, NPUModelRunner): matches the existing OmniNPUModelRunner shape, but pulls in execute_model / sample_tokens / _capture_talker_mtp_graphs / capture_model (all device-specific) which NPU still has to override — the surface saved isn't worth the silent-regression risk (any new helper added to GPU runner that touches torch.cuda.* would auto-leak to NPU).

   Recommend mixin first; revisit sibling-merge if 3+ shared helpers accumulate without device-specific contamination.